#!/usr/bin/env python3
# -*- encoding: utf-8 -*-
'''
@File: install.py
@Description: Nuclei2Snort 项目安装和配置脚本
@Author: nuclei2snort
@Version: 1.0.0
'''

import os
import sys
import subprocess
from pathlib import Path


def print_banner():
    """打印欢迎横幅"""
    banner = """
    ╔═══════════════════════════════════════════════════════════════╗
    ║                     Nuclei2Snort 安装程序                     ║
    ║               将Nuclei POC转换为Snort规则的工具               ║
    ╚═══════════════════════════════════════════════════════════════╝
    """
    print(banner)


def check_python_version():
    """检查Python版本"""
    print("🔍 检查Python版本...")
    if sys.version_info < (3, 7):
        print("❌ 错误：需要Python 3.7或更高版本")
        print(f"   当前版本：{sys.version}")
        sys.exit(1)
    print(f"✅ Python版本检查通过：{sys.version.split()[0]}")


def install_requirements():
    """安装Python依赖"""
    print("\n📦 安装Python依赖包...")
    requirements_file = Path("requirements.txt")
    
    if not requirements_file.exists():
        print("❌ 错误：未找到requirements.txt文件")
        return False
    
    try:
        cmd = [sys.executable, "-m", "pip", "install", "-r", "requirements.txt"]
        result = subprocess.run(cmd, capture_output=True, text=True)
        
        if result.returncode == 0:
            print("✅ 依赖包安装成功")
            return True
        else:
            print(f"❌ 依赖包安装失败：{result.stderr}")
            return False
    except Exception as e:
        print(f"❌ 安装过程中发生错误：{e}")
        return False


def create_directories():
    """创建必要的目录"""
    print("\n📁 创建项目目录...")
    
    directories = ['output', 'logs', 'input/pocs', 'input/cve_fail']
    
    for directory in directories:
        dir_path = Path(directory)
        if not dir_path.exists():
            dir_path.mkdir(parents=True, exist_ok=True)
            print(f"✅ 创建目录：{directory}")
        else:
            print(f"ℹ️  目录已存在：{directory}")


def setup_environment():
    """设置环境配置"""
    print("\n⚙️  配置环境变量...")
    
    config_example = Path("config.env.example")
    env_file = Path(".env")
    
    if not env_file.exists() and config_example.exists():
        print("📝 发现环境变量配置模板文件")
        
        # 询问用户是否要配置API密钥
        while True:
            choice = input("是否要配置腾讯云翻译API密钥？(y/n): ").lower().strip()
            if choice in ['y', 'yes', '是']:
                setup_api_keys()
                break
            elif choice in ['n', 'no', '否']:
                print("⏭️  跳过API密钥配置（翻译功能将不可用）")
                break
            else:
                print("请输入 y 或 n")
    else:
        print("ℹ️  环境变量配置文件已存在或模板文件不存在")


def setup_api_keys():
    """配置API密钥"""
    print("\n🔑 配置腾讯云API密钥...")
    print("请访问 https://console.cloud.tencent.com/cam/capi 获取您的API密钥")
    
    secret_id = input("请输入您的腾讯云SecretId: ").strip()
    secret_key = input("请输入您的腾讯云SecretKey: ").strip()
    
    if secret_id and secret_key:
        env_content = f"""# 腾讯云翻译API配置
TENCENT_SECRET_ID={secret_id}
TENCENT_SECRET_KEY={secret_key}

# 此文件包含敏感信息，请勿提交到代码仓库
"""
        
        with open(".env", "w", encoding="utf-8") as f:
            f.write(env_content)
        
        print("✅ API密钥配置完成")
        print("⚠️  请确保将.env文件添加到.gitignore中，避免密钥泄露")
    else:
        print("⏭️  跳过API密钥配置")


def create_gitignore():
    """创建或更新.gitignore文件"""
    print("\n🚫 更新.gitignore文件...")
    
    gitignore_entries = [
        "# 环境变量文件",
        ".env",
        "",
        "# Python缓存文件", 
        "__pycache__/",
        "*.pyc",
        "*.pyo",
        "*.pyd",
        ".Python",
        "",
        "# 日志文件",
        "logs/*.log",
        "",
        "# 临时文件",
        "*.tmp",
        "*.temp",
        "",
        "# IDE文件",
        ".vscode/",
        ".idea/",
        "*.swp",
        "*.swo",
        "",
        "# 输出文件",
        "output/*.rules"
    ]
    
    gitignore_path = Path(".gitignore")
    existing_content = ""
    
    if gitignore_path.exists():
        with open(gitignore_path, "r", encoding="utf-8") as f:
            existing_content = f.read()
    
    # 添加缺失的条目
    new_entries = []
    for entry in gitignore_entries:
        if entry and entry not in existing_content:
            new_entries.append(entry)
    
    if new_entries:
        with open(gitignore_path, "a", encoding="utf-8") as f:
            if existing_content and not existing_content.endswith('\n'):
                f.write('\n')
            f.write('\n'.join(new_entries) + '\n')
        print("✅ .gitignore文件已更新")
    else:
        print("ℹ️  .gitignore文件无需更新")


def run_test():
    """运行测试"""
    print("\n🧪 运行功能测试...")
    
    example_file = Path("example_poc.yaml")
    if example_file.exists():
        try:
            cmd = [sys.executable, "run.py", "--file", str(example_file)]
            result = subprocess.run(cmd, capture_output=True, text=True, timeout=30)
            
            if result.returncode == 0:
                print("✅ 功能测试通过")
                return True
            else:
                print(f"⚠️  测试警告：{result.stderr}")
                return True  # 可能因为API密钥未配置导致的警告，但不影响基本功能
        except subprocess.TimeoutExpired:
            print("⚠️  测试超时，可能因为网络问题")
            return True
        except Exception as e:
            print(f"❌ 测试失败：{e}")
            return False
    else:
        print("⏭️  跳过测试（示例文件不存在）")
        return True


def print_completion_info():
    """打印安装完成信息"""
    print("\n" + "="*60)
    print("🎉 安装完成！")
    print("="*60)
    
    print("\n📖 使用说明：")
    print("1. 查看帮助信息：")
    print("   python run.py --help")
    
    print("\n2. 转换单个文件：")
    print("   python run.py --file your_poc.yaml")
    
    print("\n3. 批量转换目录：")
    print("   python run.py --dir input/pocs/")
    
    print("\n4. 使用并发处理：")
    print("   python run.py --dir input/pocs/ --workers 8")
    
    print("\n📁 项目目录结构：")
    print("├── input/     # 输入POC文件目录")
    print("├── output/    # 输出Snort规则目录")
    print("├── logs/      # 日志文件目录")
    print("└── .env       # 环境变量配置文件（如已配置）")
    
    print("\n🔗 更多信息请查看 README.md")
    print("\n感谢使用 Nuclei2Snort！")


def main():
    """主安装函数"""
    print_banner()
    
    try:
        # 检查Python版本
        check_python_version()
        
        # 安装依赖
        if not install_requirements():
            print("\n❌ 安装失败：依赖包安装失败")
            sys.exit(1)
        
        # 创建目录
        create_directories()
        
        # 设置环境
        setup_environment()
        
        # 创建.gitignore
        create_gitignore()
        
        # 运行测试
        run_test()
        
        # 显示完成信息
        print_completion_info()
        
    except KeyboardInterrupt:
        print("\n\n⚠️  安装被用户中断")
        sys.exit(1)
    except Exception as e:
        print(f"\n❌ 安装过程中发生错误：{e}")
        sys.exit(1)


if __name__ == "__main__":
    main() 